#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : xfyun.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/8 10:28
import asyncio
from typing import List

from apps.model_utils.base import BaseModelUtil
from apps.utils import async_http_client


class Xfyun(BaseModelUtil):

    def __init__(self, api_key, prompt: str = None, message: List[dict] = None,
                 temperature: float = 0.9, top_p: float = 1, max_tokens=4096, top_k: int = 4,
                 presence_penalty=0, frequency_penalty=0, *args, **kwargs):
        self.usage_tokens = 0
        self.api_key = api_key
        self.prompt = prompt
        self.message = message if message else []
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.top_k = top_k
        self.presence_penalty = presence_penalty
        self.frequency_penalty = frequency_penalty

    async def chat(self, query, base_url, model_name, **kwargs):
        message = self.get_message()
        print(message, message)
        message.append({"role": "user", "content": query})
        if self.prompt:
            message.insert(0, {"role": "system", "content": self.prompt})
        print({
            'messages': message,
            'model': model_name,
            'temperature': self.temperature,
            'top_p': self.top_p,
            'max_tokens': self.max_tokens,
            'top_k': self.top_k,
            'presence_penalty': self.presence_penalty,
            'frequency_penalty': self.frequency_penalty,
        })
        response = await async_http_client.post(
            url=base_url, json={
                'messages': message,
                'model': model_name,
                'temperature': self.temperature,
                'top_p': self.top_p,
                'max_tokens': self.max_tokens,
                'top_k': self.top_k,
                'presence_penalty': self.presence_penalty,
                'frequency_penalty': self.frequency_penalty,
                "stream": False
            },
            headers=self._header()
        )
        print(response)
        self.usage_tokens = response.get("usage", {}).get('total_tokens')
        return response.get("choices", [{}])[0].get("message", {}).get("content", "")

    def _header(self) -> dict:
        return {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}

    def get_message(self):
        if self.message:
            _ = []
            for i in self.message:
                _.append({"role": "user", "content": i.get('query')})
                _.append({"role": "assistant", "content": i.get('answer')})
            return _
        else:
            return []

    async def get_usage_tokens(self):
        return self.usage_tokens
